# -----------------------------------------------------------------------------
# Copyright (c) 2025, Southeast University (China)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name of the copyright holders nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Jiajie Xu, Yanfang Zhang, Jiaqi Gao, Leyun Tian
# -----------------------------------------------------------------------------


from Data_var import *
import work_var
import re
from multiprocessing import Process, Manager, Pool
import os
import pickle
import bisect
import warnings


class TimingLib:
    """TimingLibs for current technology"""

    """
    self.vth_corner -> (vth, corner, voltage, temperature) of libs
    self.NLDM -> a dict for cell NLDM models: self.vth_corner -> { cell name (ND2D0BWP16P90) -> NLDM r/f Delay/Trans Lookup Tables }
    self._load_liberty_(self) -> load a pre-saved liberty file
    self.NLDM_Timing(self, corner, cellname, arc, in_rf, sense_unate, trans, caps): calculate delay and trans based on the NLDM model.
        Args: corner: (corner, voltage, temperature), cellname: cell name, arc: the timing arc to be calculated , in_rf: input transition direction, sense_unate: positive/negative unate of timing sense, trans: input transition time (index1), capacitance: output capacitance (index2)
        returns delay or trans for the cell, dict(delay/trans) : {r/f: value}, possibly two values for cell like XOR2D0P75BWP16P90
    self.interpolation(self, trans, caps): Using equation AX + BY + CXY + D, to interpolate based on the two dimension LUT
    """

    def __init__(self):
        # Get lib names under work_dir/lib_dir
        lib_files = [f for f in os.listdir(work_var.lib_dir) if f.endswith(".lib")]
        pattern = re.compile(
            r"tcbn12ffcllbwp16p90(?P<vth>lvt|ulvt|)(?P<corner>ssgnp|ffgnp|tt)(?P<voltage>[^v]+)v(?P<temperature>[^c]+)c_ccs.lib"
        )
        self.vth_corner = {}  # (vth, corner, voltage, temperature) -> lib name
        self.NLDM = {}  # (corner, voltage, temperature) -> cells
        vth_corner_libs = (
            {}
        )  # (vth, corner, voltage, temperature) -> cells , for merging to NLDM

        self.footprints_to_cells = (
            []
        )  # footprint -> cellname of different vth, for gate sizing only

        for lib in lib_files:
            match = pattern.match(lib)
            if match:
                vth_value = match.group("vth") if match.group("vth") else "svt"
                self.vth_corner[
                    (
                        vth_value,
                        match.group("corner"),
                        match.group("voltage"),
                        match.group("temperature"),
                    )
                ] = lib
            else:
                raise ValueError(
                    "lib name not match! please use the libs given by factory!"
                )

        # check saved liberty
        os.makedirs(Liberty_restore_path, exist_ok=True)
        saved_liberty = [f.split(".")[0] for f in os.listdir(Liberty_restore_path)]
        unsaved_vth_corner = {}
        saved_vth_corner = {}
        for vth_corner in self.vth_corner.keys():
            if "_".join(vth_corner) not in saved_liberty:
                unsaved_vth_corner[vth_corner] = self.vth_corner[vth_corner]
            else:
                saved_vth_corner[vth_corner] = self.vth_corner[vth_corner]

        max_processes = 8
        # Reading saved liberty with pickle
        if saved_vth_corner.keys():
            print(
                f"Loading saved design liberty files with maximum {max_processes} processes."
            )
        with Pool(processes=max_processes) as pool:
            results = pool.map(
                self._load_liberty_,
                [(key, Liberty_restore_path) for key in saved_vth_corner.keys()],
            )
        vth_corner_libs.update(results)

        processes = []
        manager = Manager()
        results = manager.list([None] * len(unsaved_vth_corner.values()))
        if unsaved_vth_corner.keys():
            print(
                f"Reading unsaved design liberty files with maximum {max_processes} processes."
            )

        try:
            for i, lib in enumerate(unsaved_vth_corner.values()):
                if len(processes) >= max_processes:
                    processes[0].join()  # wait for the earliest process
                    processes.pop(0)  # remove complete process
                process = Process(
                    target=TimingLib_Parser.Read_TimingLib,
                    args=(os.path.join(work_var.lib_dir, lib), i, results),
                )
                processes.append(process)
                process.start()

            for process in processes:
                process.join()  # wait for all process complete

        except Exception as e:
            print(f"Error occurred: {e}")
            raise

        # Save liberty for future utilize
        for vth_corner, result in zip(unsaved_vth_corner.keys(), results):
            save_path = os.path.join(Liberty_restore_path, "_".join(vth_corner))
            vth_corner_libs[vth_corner] = result
            with open(f"{save_path}.pkl", "wb") as f:
                pickle.dump(result, f)

        vth = []

        for key, value in vth_corner_libs.items():
            if (key[1], key[2], key[3]) not in self.NLDM.keys():
                self.NLDM[(key[1], key[2], key[3])] = value
            else:
                self.NLDM[(key[1], key[2], key[3])].update(value)

            # footprint mapped to cellname for each vth for gate sizing
            if key[0] not in vth:
                vth.append(key[0])
                footprint_to_cell = {}
                for cell in vth_corner_libs[key].values():
                    if cell.footprint not in footprint_to_cell.keys():
                        footprint_to_cell[cell.footprint] = [cell.name]
                    else:
                        footprint_to_cell[cell.footprint].append(cell.name)
                    self.footprints_to_cells.append(footprint_to_cell)

            print(f"{self.vth_corner[key]} Readed.")

        print("\rAll liberty readed!")

    def _load_liberty_(self, args):
        """Loading saved liberties"""
        key, Liberty_restore_path = args
        save_path = os.path.join(Liberty_restore_path, "_".join(key))
        with open(f"{save_path}.pkl", "rb") as f:
            result = pickle.load(f)
        return key, result

    def _interpolation_(self, LUT, trans, caps):
        # Determine whether to use 'delay' or 'trans' from LUT
        if hasattr(LUT, "delay"):
            values = LUT.delay
        elif hasattr(LUT, "trans"):
            values = LUT.trans
        elif hasattr(LUT, "constrain"):
            values = LUT.constrain
        else:
            raise AttributeError(
                "LUT must contain either 'delay'/'trans'/'constrain' attribute"
            )

        index1 = LUT.index1
        index2 = LUT.index2
        # Check for valid index lengths
        if len(index1) < 2 or len(index2) < 2:
            raise ValueError("Both index arrays must contain at least two elements")

        # Find interpolation positions for trans
        i = bisect.bisect_left(index1, trans)
        if i == 0:
            x_low, x_high = 0, 1
        elif i == len(index1):
            x_low, x_high = len(index1) - 2, len(index1) - 1
        else:
            x_low, x_high = i - 1, i

        # Find interpolation positions for caps
        j = bisect.bisect_left(index2, caps)
        if j == 0:
            y_low, y_high = 0, 1
        elif j == len(index2):
            y_low, y_high = len(index2) - 2, len(index2) - 1
        else:
            y_low, y_high = j - 1, j

        # Get boundary values
        x0, x1 = index1[x_low], index1[x_high]
        y0, y1 = index2[y_low], index2[y_high]

        # Get corresponding matrix values
        try:
            z00 = values[x_low][y_low]
            z01 = values[x_low][y_high]
            z10 = values[x_high][y_low]
            z11 = values[x_high][y_high]
        except IndexError:
            raise IndexError("Value matrix dimensions do not match index arrays")

            # Calculate interpolation coefficients
        try:
            C = (z11 - z10 - z01 + z00) / ((x1 - x0) * (y1 - y0))
            B = (z01 - z00) / (y1 - y0) - C * x0
            A = (z10 - z00) / (x1 - x0) - C * y0
            D = z00 - A * x0 - B * y0 - C * x0 * y0
        except ZeroDivisionError:
            raise ValueError("Index values must be strictly increasing")

        # Calculate final result
        result = A * trans + B * caps + C * trans * caps + D

        # Check for boundary warnings
        min1, max1 = min(index1), max(index1)
        if trans < min1 - 0.1 * min1:
            warnings.warn(
                f"trans value {trans} exceeds lower bound by more than 10%", UserWarning
            )
        elif trans > max1 + 0.1 * max1:
            warnings.warn(
                f"trans value {trans} exceeds upper bound by more than 10%", UserWarning
            )

        min2, max2 = min(index2), max(index2)
        if caps < min2 - 0.1 * min2:
            warnings.warn(
                f"caps value {caps} exceeds lower bound by more than 10%", UserWarning
            )
        elif caps > max2 + 0.1 * max2:
            warnings.warn(
                f"caps value {caps} exceeds upper bound by more than 10%", UserWarning
            )

        return result

    def NLDM_Timing(
        self,
        corner,
        cellname,
        arc,
        out_rf,
        sense_unate,
        timing_type,
        when,
        sdf_cond,
        trans,
        caps,
    ):
        cell = self.NLDM[corner][cellname]
        # if sense_unate == 'positive_unate':
        #     out_rf = in_rf
        # elif sense_unate == 'negative_unate':
        #     out_rf = 'r' if in_rf == 'f' else 'f'
        # if the key exist, use interpolation to calculate delay and trans
        NLDM_delay = cell.delay.get(
            (arc, out_rf, sense_unate, timing_type, when, sdf_cond), None
        )
        NLDM_trans = cell.trans.get(
            (arc, out_rf, sense_unate, timing_type, when, sdf_cond), None
        )
        if NLDM_trans:
            delay = self._interpolation_(NLDM_delay, trans, caps)
            if NLDM_trans.isscalar:
                out_trans = 0
            else:
                out_trans = self._interpolation_(NLDM_trans, trans, caps)
        else:
            # else, return None
            delay = None
            out_trans = None
        return delay, out_trans

    def NLDM_Constraint(
        self,
        corner,
        cellname,
        arc,
        out_rf,
        sense_unate,
        timing_type,
        when,
        sdf_cond,
        ftrans,
        ttrans,
    ):
        cell = self.NLDM[corner][cellname]
        NLDM_cons = cell.constrain.get(
            (arc, out_rf, sense_unate, timing_type, when, sdf_cond), None
        )
        cons = self._interpolation_(NLDM_cons, ftrans, ttrans)

        return cons

    def Get_Candidate(self, cellname):
        footprint = next(iter(self.NLDM.values()))[cellname].footprint
        for item in self.footprints_to_cells:
            if cellname in item[footprint]:
                return item[footprint]

    def print_lib(self):
        """
        print the timing lib information saved
        """
        for key, value in self.NLDM.items():
            print(
                f"------------------------------------------ Lib Information for corner:{key[0]} voltage:{key[1]} temperature{key[2]} ------------------------------------------"
            )
            for cell in value.values():
                print(cell)


if __name__ == "__main__":
    TL = TimingLib()
    print(
        TL.NLDM_Timing(
            ("ssgnp", "0p63", "125"),
            "OAI211D0BWP16P90",
            "B->ZN",
            "r",
            "negative_unate",
            "combinational",
            "A1&A2&C",
            "A1==1'b1&&A2==1'b1&&C==1'b1",
            0.048682,
            0.000705,
        )
    )
    print(TL.Get_Candidate("OAI211D0BWP16P90"))
